package org.ws.task.schedule;

import com.google.common.collect.Lists;
import com.google.common.collect.Maps;
import org.apache.commons.collections4.CollectionUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.ws.task.item.ExecuteContext;
import org.ws.task.item.ItemLeaderProvider;
import org.ws.task.item.TaskInstance;
import org.ws.task.repository.LeaderContextRepository;
import org.ws.task.repository.TaskRepository;
import org.ws.task.schedule.dispatcher.Dispatcher;
import org.ws.task.schedule.dispatcher.DispatcherFactory;
import org.ws.task.model.protocol.PullItemRequest;
import org.ws.task.model.TaskConfig;
import org.ws.task.model.TaskItem;

import java.util.Collection;
import java.util.Collections;
import java.util.List;
import java.util.Map;
import java.util.concurrent.atomic.AtomicReference;

public class DefaultItemSchedule<T> implements ItemSchedule<T> {

    private static Logger log = LoggerFactory.getLogger(DefaultItemSchedule.class.getName());

    private final Map<String,AtomicReference<ScheduleStatus>> scheduleStatusMap = Maps.newConcurrentMap();

    @Override
    public void schedule(String taskName, String nodeId, int size,
                         ScheduleCallback<TaskItem<T>> callback) {
        TaskInstance taskInstance = TaskRepository.getInstance().getTaskInstance(taskName);
        Collection<TaskItem<T>> taskItems = doSchedule(taskInstance, nodeId, size);
        if(CollectionUtils.isNotEmpty(taskItems)){
            for (TaskItem<T> item : taskItems) {
                try {
                    ScheduleStatus scheduleStatus = callback.callback(item);
                    this.getScheduleStatus(taskName,nodeId).set(scheduleStatus);
                } catch (Exception e) {
                    log.error(e.getMessage(),e);
                    this.getScheduleStatus(taskName,nodeId).set(ScheduleStatus.STOP);
                }
            }
        }
        else {
            this.getScheduleStatus(taskName,nodeId).set(ScheduleStatus.STOP);
        }
    }

    private Collection<TaskItem<T>> doSchedule(TaskInstance taskInstance, String nodeId, int size){
        String taskName = taskInstance.getTaskConfig().getName();
        ExecuteContext context = LeaderContextRepository.getInstance().getContext(taskName, nodeId);
        // context为空加载
        if(context.isEmpty()){
            Collection<TaskItem<T>> taskItems = loadByProvider(taskInstance, nodeId, size);
            if(CollectionUtils.isNotEmpty(taskItems)){
                TaskConfig.DispatchStrategy dispatchStrategy = taskInstance.getTaskConfig().getDispatchStrategy();
                TaskConfig.SendModel sendModel = taskInstance.getTaskConfig().getSendModel();
                Dispatcher dispatcher = DispatcherFactory.getInstance().getDispatcher(dispatchStrategy,sendModel);
                for (TaskItem<T> item : taskItems) {
                    dispatcher.dispatch(item,taskName);
                }
            }
            else {
                context.printSize();
            }
        }
        List<TaskItem<T>> list = Lists.newArrayList();
        for (int i = 0; i < size; i++) {
            if(!context.isEmpty()){
                list.add(context.poll());
            }
        }
        return list;
    }

    @Override
    public void ready(String taskName,String nodeId) {
        if(this.getScheduleStatus(taskName,nodeId).get()==ScheduleStatus.STOP) {
            this.getScheduleStatus(taskName, nodeId).set(ScheduleStatus.READY);
        }
    }

    /**
     * pull:主动拉取数据
     * push：事件驱动
     * @param
     */
    protected Collection<TaskItem<T>> loadByProvider(TaskInstance taskInstance, String nodeId, int size){
        TaskConfig taskConfig = taskInstance.getTaskConfig();
        if(taskConfig.getTaskType() == TaskConfig.TaskType.PULL){
            ItemLeaderProvider itemLeaderProvider = taskInstance.getItemLeaderProvider();
            PullItemRequest pullItemRequest = new PullItemRequest();
            pullItemRequest.setTaskName(taskConfig.getName());
            pullItemRequest.setSize(size);
            pullItemRequest.setNodeId(nodeId);
            pullItemRequest.setTimeMillis(System.currentTimeMillis());

            return itemLeaderProvider.getItems(pullItemRequest);
        }
        return Collections.emptyList();
    }

    @Override
    public void shutdown(String taskName,String nodeId) {
        this.getScheduleStatus(taskName,nodeId).set(ScheduleStatus.FINISHED);
    }

    @Override
    public boolean isContinue(String taskName,String nodeId) {
        return this.getScheduleStatus(taskName,nodeId).get().isContinuable();
    }

    private AtomicReference<ScheduleStatus> getScheduleStatus(String taskName, String nodeId){
        String key = taskName+nodeId;
        if(!this.scheduleStatusMap.containsKey(key)){
            this.scheduleStatusMap.put(key,new AtomicReference<>(ScheduleStatus.READY));
        }
        return this.scheduleStatusMap.get(key);
    }
}
